#include "rtda/heap/class_loader.h"
#include "rtda/heap/class.h"
#include "classfile/class_file.hpp"
#include <iostream>

namespace jvm {
namespace heap {

ClassLoader::ClassLoader(std::shared_ptr<Classpath> pClasspath) 
    : m_pClasspath(pClasspath) {}

std::shared_ptr<Class> ClassLoader::loadClass(const std::string& name) {
    auto it = m_classMap.find(name);
    if (it != m_classMap.end()) {
        return it->second;
    }
    return loadNonArrayClass(name);
}

std::shared_ptr<Class> ClassLoader::loadNonArrayClass(const std::string& name) {
    auto [data, entry] = readClass(name);
    std::shared_ptr<Class> pClass = defineClass(data);
    link(pClass);
    std::cout << "[Loaded " << name << " from " << entry.toString() << "]\n";
    return pClass;
}

std::pair<std::vector<uint8_t>, classpath::Entry> ClassLoader::readClass(const std::string& name) {
    try {
        return m_pClasspath->read_class(name);
    } catch (const std::exception& e) {
        throw std::runtime_error("java.lang.ClassNotFoundException: " + name);
    }
}

std::shared_ptr<Class> ClassLoader::defineClass(const std::vector<uint8_t>& data) {
    std::shared_ptr<Class> pClass = parseClass(data);
    pClass->m_pLoader = this;
    resolveSuperClass(pClass);
    resolveInterfaces(pClass);
    m_classMap[pClass->m_name] = pClass;
    return pClass;
}

std::shared_ptr<Class> ClassLoader::parseClass(const std::vector<uint8_t>& data) {
    try {
        std::shared_ptr<classfile::ClassFile> pClassFile = classfile::parse(data);
        return new Class(pClassFile);
    } catch (const std::exception& e) {
        throw std::runtime_error("java.lang.ClassFormatError");
    }
}

void ClassLoader::resolveSuperClass(std::shared_ptr<Class> pClass) {
    if (pClass->m_name != "java/lang/Object") {
        pClass->m_pSuperClass = pClass->m_pLoader->loadClass(pClass->m_superClassName);
    }
}

void ClassLoader::resolveInterfaces(std::shared_ptr<Class> pClass) {
    if (!pClass->m_interfaceNames.empty()) {
        pClass->m_interfaces.resize(pClass->m_interfaceNames.size());
        for (size_t i = 0; i < pClass->m_interfaceNames.size(); ++i) {
            pClass->m_interfaces[i] = pClass->m_pLoader->loadClass(pClass->m_interfaceNames[i]);
        }
    }
}

void ClassLoader::link(std::shared_ptr<Class> pClass) {
    verify(pClass);
    prepare(pClass);
}

// void ClassLoader::verify(std::shared_ptr<Class> pClass) {
void verify(std::shared_ptr<Class> pClass) {
    // TODO: Implement verification
}

// void ClassLoader::prepare(std::shared_ptr<Class> pClass) {
void prepare(std::shared_ptr<Class> pClass) {
    calcInstanceFieldSlotIds(pClass);
    calcStaticFieldSlotIds(pClass);
    allocAndInitStaticVars(pClass);
}

// void ClassLoader::calcInstanceFieldSlotIds(std::shared_ptr<Class> pClass) {
void calcInstanceFieldSlotIds(std::shared_ptr<Class> pClass) {
    uint slotId = 0;
    if (pClass->m_pSuperClass != nullptr) {
        slotId = pClass->m_pSuperClass->m_instanceSlotCount;
    }
    
    for (std::shared_ptr<Field> pField : pClass->m_fields) {
        if (!pField->isStatic()) {
            pField->setSlotId(slotId);
            slotId++;
            if (pField->isLongOrDouble()) {
                slotId++;
            }
        }
    }
    pClass->m_instanceSlotCount = slotId;
}

// void ClassLoader::calcStaticFieldSlotIds(std::shared_ptr<Class> pClass) {
void calcStaticFieldSlotIds(std::shared_ptr<Class> pClass) {
    uint slotId = 0;
    for (std::shared_ptr<Field> pField : pClass->m_fields) {
        if (pField->isStatic()) {
            pField->setSlotId(slotId);
            slotId++;
            if (pField->isLongOrDouble()) {
                slotId++;
            }
        }
    }
    pClass->m_staticSlotCount = slotId;
}

// void ClassLoader::allocAndInitStaticVars(std::shared_ptr<Class> pClass) {
void allocAndInitStaticVars(std::shared_ptr<Class> pClass) {
    pClass->m_pStaticVars = newSlots(pClass->m_staticSlotCount);
    for (std::shared_ptr<Field> pField : pClass->m_fields) {
        if (pField->isStatic() && pField->isFinal()) {
            initStaticFinalVar(pClass, pField);
        }
    }
}

// void ClassLoader::initStaticFinalVar(std::shared_ptr<Class> pClass, std::shared_ptr<Field> pField) {
void initStaticFinalVar(std::shared_ptr<Class> pClass, std::shared_ptr<Field> pField) {
    std::shared_ptr<Slots> pVars = pClass->getStaticVars();
    if (pVars == nullptr) {
        throw std::runtime_error("Static variables not allocated");
    }
    std::shared_ptr<ConstantPool> pCp = pClass->getConstantPool();
    if (pCp == nullptr) {
        throw std::runtime_error("Constant pool not initialized");
    }
    uint16_t cpIndex = pField->getConstValueIndex();
    uint slotId = pField->getSlotId();

    if (cpIndex > 0) {
        std::string descriptor = pField->getDescriptor();
        if (descriptor == "Z" || descriptor == "B" || descriptor == "C" || descriptor == "S" || descriptor == "I") {
            int32_t val = pCp->getConstant(cpIndex).getInt();
            pVars->setInt(slotId, val);
        } else if (descriptor == "J") {
            int64_t val = pCp->getConstant(cpIndex).getLong();
            pVars->setLong(slotId, val);
        } else if (descriptor == "F") {
            float val = pCp->getConstant(cpIndex).getFloat();
            pVars->setFloat(slotId, val);
        } else if (descriptor == "D") {
            double val = pCp->getConstant(cpIndex).getDouble();
            pVars->setDouble(slotId, val);
        } else if (descriptor == "Ljava/lang/String;") {
            throw std::runtime_error("String constant initialization not implemented");
        }
    }
}

} // namespace heap
} // namespace jvm